import audioop
import os
import shutil
import time
import wave
from glob import glob

from moviepy.audio.io.AudioFileClip import AudioFileClip
from tqdm import tqdm
import numpy as np

from config import Config
from models import WakeupModule

CHUNK = Config["chunk"]
CPS = Config["chunks_per_sample"]
Sample_Len = CHUNK * CPS

filters = ["/*.wav",  # 0 all

           "/负样本*.wav",  # 1 负样本.mp3
           "/* *.wav",  # 2 人工语音.zip 和 自采集
           "/G*.wav",  # 3 购买siri.zip 和 hey_siri.zip
           "/20210531_*.wav",  # 4 赵佳朋 siri.rar
           "/20210607_*.wav",  # 5 2021-6-7.rar 和 Siri语音数据.zip
           "/20220316_*.wav",  # 6 uploads.zip 和 audio.zip

           "/record_*.wav",  # 7 自采集 或丢失时间戳
           "/old_*.wav",  # 8 以前裁剪好的样本，但是丢失原始语料
           "/crop_*.wav",  # 9 以前裁剪好的样本，有原始语料包，但是裁剪方法或有不同
           ]


def positive_wav_data_(src, des):
    with wave.open(src, 'r') as wf:
        _params = wf.getparams()
        assert (_params.sampwidth == 2)
        assert (_params.nchannels == 1)
        assert (_params.framerate == Config["rate"])

        cnt = wf.getnframes() // CHUNK
        wav_data = []
        rms_list = [0 for _ in range(cnt)]
        for i in range(cnt):
            wf.setpos(CHUNK * i)
            data = wf.readframes(CHUNK)
            wav_data.append(data)
            rms_list[i] = audioop.rms(data, 2)

    max_start = 0
    for i in range(cnt):
        if rms_list[i] > 500:
            max_start = i
            break

    # assert(max_start > 0)

    max_end = cnt - 1
    for i in range(cnt - 1, -1, -1):
        if rms_list[i] > 300:
            max_end = i
            break

    # assert(max_end < cnt-1)

    start = max_start
    for i in range(max_start - 1, -1, -1):
        if rms_list[i] < 200:
            start = i
            break

    assert (max_start > start and max_start - start < 50)

    end = max_end
    for i in range(max_end, cnt):
        if rms_list[i] < 100:
            end = i
            break

    if end <= max_end or end - max_end >= 50:
        print(end - max_end)

    # assert (end > max_end and end - max_end < 50)

    with wave.open(des, 'w') as wf:
        params = (1, 2, Config["rate"], 0, 'NONE', 'Uncompressed')
        wf.setparams(params)
        wf.writeframes(wav_data[start:end])


def wave_to_uniwave_(src_path, des_path):
    try:
        wf = wave.open(src_path, 'r')
        _params = wf.getparams()
        data = wf.readframes(_params.nframes)
        if _params.nchannels != 1:
            assert (_params.nchannels == 2)
            data = audioop.tomono(data, 2, 0.5, 0.5)
        if _params.sampwidth != 2:
            assert (_params.sampwidth == 1)
            data = audioop.lin2lin(data, _params.sampwidth, 2)
            data = audioop.bias(data, 1, 128)
        if _params.framerate != Config["rate"]:
            data, _ = audioop.ratecv(data, 2, 1, _params.framerate, Config["rate"], None)
        with wave.open(des_path, 'w') as out_f:
            out_f.setparams((1, 2, Config["rate"], 0, 'NONE', 'Uncompressed'))
            out_f.writeframes(data)
        wf.close()
    except Exception as e:
        print(e)
        print(src_path)


def mp3_to_uniwave(root):
    files = glob("%s/*.mp3" % root)
    for file in files:
        name = os.path.basename(file)
        (shot_name, extension) = os.path.splitext(name)
        des_path = os.path.join("../data/temp", shot_name + ".wav")
        audio_clip = AudioFileClip(file, fps=Config["rate"])
        audio_clip.write_audiofile(filename=des_path, fps=Config["rate"], nbytes=2, codec='pcm_s16le')
        wave_to_uniwave_(des_path, des_path)


def wave_to_uniwave(root):
    files = glob("%s/*.wav" % root)
    for file in tqdm(files):
        name = os.path.basename(file)
        (shot_name, extension) = os.path.splitext(name)
        des_path = os.path.join("../data/temp", shot_name + ".wav")
        wave_to_uniwave_(file, des_path)


def rename_wave(root, prefix):
    files = sorted(glob("%s/*.wav" % root), key=os.path.getmtime)
    for i in range(len(files)):
        f = files[i]
        root_dir = os.path.dirname(f)
        basename = os.path.basename(f)
        short_name, short_ext = os.path.splitext(basename)
        des = os.path.join(root_dir, prefix + "%05d%s" % (i, short_ext))
        os.rename(f, des)


def rename_by_time(root, flt):
    files = sorted(glob(root + flt), key=os.path.getmtime)
    for file in files:
        mtm = os.path.getmtime(file)
        tm = time.localtime(mtm)
        name = "%04d%02d%02d %02d%02d%02d.wav" % (
            tm.tm_year, tm.tm_mon, tm.tm_mday, tm.tm_hour, tm.tm_min, tm.tm_sec)
        des = os.path.join(root, name)
        if os.path.exists(des):
            print(file, des)
        else:
            os.rename(file, des)


def rename_by_filter(root, flt, prefix):
    files = glob(root + "/**" + flt, recursive=True)
    files = sorted(files, key=os.path.getmtime)
    for i in tqdm(range(len(files))):
        f = files[i]
        root_dir = os.path.dirname(f)
        basename = os.path.basename(f)
        short_name, short_ext = os.path.splitext(basename)
        des = os.path.join(root_dir, prefix + "%05d%s" % (i, short_ext))
        os.rename(f, des)


def check_by_time(root, flt, moveto=None):
    files = glob(root + flt)
    for file in files:
        mtm = os.path.getmtime(file)
        tm = time.localtime(mtm)
        # if tm.tm_year == 2021 and tm.tm_mon == 4 and tm.tm_mday == 23 and tm.tm_hour == 19:
        if tm.tm_year == 2021 and tm.tm_mon == 5 and (tm.tm_mday == 14 or tm.tm_mday == 15):
            print(file)
            if moveto:
                name = os.path.basename(file)
                des = os.path.join(moveto, name)
                shutil.move(file, des)


def classify_uniwave(root, moveto=None):
    def recognize_uniwave():
        with wave.open(file, 'r') as wf:
            cnt = wf.getnframes() // CHUNK
            cnt -= CPS
            for i in range(0, cnt, 1):
                pos = wf.getnframes() - Sample_Len
                wf.setpos(pos - CHUNK * i)
                data = wf.readframes(Sample_Len)
                outputs = wakeup.recognize(data)
                if outputs[0] > Config["confidence"]:
                    return 1
        return 0

    wakeup = WakeupModule(Config["weight_path"]).eval()
    files = glob("%s/*.wav" % root)
    for file in files:
        ret = recognize_uniwave()
        if ret == 0:
            print(file)
            if moveto:
                name = os.path.basename(file)
                des = os.path.join(moveto, name)
                shutil.move(file, des)


def unniwave_crop_0sample(root, step=CPS // 8):
    files = glob("%s/*.wav" % root)
    for file in tqdm(files):
        name = os.path.basename(file)
        (shot_name, extension) = os.path.splitext(name)
        with wave.open(file, 'r') as wf:
            cnt = wf.getnframes() // CHUNK
            if CPS > cnt > CPS - 5:
                data = wf.readframes(CHUNK * cnt)
                data = b'\x00' * (CPS - cnt) * 2 * CHUNK + data
                des_path = os.path.join("../data/temp", shot_name + "_%02d.wav" % 0)
                with wave.open(des_path, 'w') as out_f:
                    out_f.setparams((1, 2, Config["rate"], 0, 'NONE', 'Uncompressed'))
                    out_f.writeframes(data)
            else:
                cnt -= CPS
                for i in range(0, cnt, step):
                    pos = wf.getnframes() - Sample_Len
                    wf.setpos(pos - CHUNK * i)
                    data = wf.readframes(Sample_Len)
                    des_path = os.path.join("../data/temp", shot_name + "_%02d.wav" % i)
                    with wave.open(des_path, 'w') as out_f:
                        out_f.setparams((1, 2, Config["rate"], 0, 'NONE', 'Uncompressed'))
                        out_f.writeframes(data)


def unniwave_crop_1sample(root):
    files = glob("%s/*.wav" % root)
    for file in tqdm(files):
        name = os.path.basename(file)
        des_path = os.path.join("../data/temp", name)
        positive_wav_data_(file, des_path)


def unniwave_reco_1sample(root):
    def save_best(best):
        (ii, out, loss, data) = best
        basename = os.path.basename(file)
        short_name, short_ext = os.path.splitext(basename)
        des = os.path.join("corpus/temp", "%s_%02d_%.4f%s" % (short_name, ii, loss[0], short_ext))
        wf = wave.open(des, 'w')
        params = (1, 2, Config["rate"], 0, 'NONE', 'Uncompressed')
        wf.setparams(params)
        wf.writeframes(data)
        wf.close()

    def parse_result(result):
        saved = 0
        best = None
        for i in range(len(result)):
            (ii, outputs, data) = result[i]
            if best is not None and ii - best[0] > CPS:
                save_best(best)
                saved += 1
                best = None

            if outputs[0] > Config["confidence"]:
                if best is None:
                    best = (ii, outputs, data)
                elif best is not None and outputs[0] < best[1][0]:
                    best = (ii, outputs, data)

        if best is not None:
            save_best(best)
            saved += 1

        return saved

    def make_sample_():
        result = []
        with wave.open(file, 'r') as wf:
            cnt = wf.getnframes() // CHUNK
            cnt -= CPS
            for i in range(0, cnt, 1):
                pos = wf.getnframes() - Sample_Len
                wf.setpos(pos - CHUNK * i)
                data = wf.readframes(Sample_Len)
                outputs = wakeup.recognize(data)
                result.append((i, outputs, data))
        return parse_result(result)

    wakeup = WakeupModule(Config["weight_path"]).eval()
    files = glob("%s/*.wav" % root)
    for file in tqdm(files):
        ret = make_sample_()
        if ret == 0:
            print(file)


def check_mute_sample(root, flag='all', moveto=None):
    # files = glob(root + "/**" + filters[0], recursive=True)
    files = glob(root + filters[0])
    for file in files:
        with wave.open(file, 'r') as wf:
            data = wf.readframes(wf.getnframes())
        mm = max(abs(np.frombuffer(data, dtype=np.int16))) + 1
        mmm = audioop.rms(data, 2)
        if (flag == 'all' and (mm < 1000 or mmm < 200)) or \
                (flag == 'max' and mm < 1000) or \
                (flag == 'rms' and mmm < 200):
            print(file)
            if moveto:
                name = os.path.basename(file)
                des = os.path.join(moveto, name)
                shutil.move(file, des)


def check_loud_sample(root, moveto=None):
    files = glob("%s/*.wav" % root)
    for file in files:
        with wave.open(file, 'r') as wf:
            data = wf.readframes(wf.getnframes())
        mmm = audioop.rms(data, 2)
        if mmm > 15000:
            print(file)
            if moveto:
                name = os.path.basename(file)
                des = os.path.join(moveto, name)
                shutil.move(file, des)


def normalize_sample(root, des_root):
    files = glob("%s/*.wav" % root)
    for file in files:
        name = os.path.basename(file)
        des_path = os.path.join(des_root, name)
        with wave.open(file, 'r') as wf:
            data = wf.readframes(wf.getnframes())
        wav_data = np.frombuffer(data, dtype=np.int16)
        wav_data = np.float32(wav_data)
        mm = max(abs(wav_data)) + 1
        wav_data = wav_data * 32000 / mm
        wav_data = np.int16(wav_data)
        data = wav_data.tobytes()
        with wave.open(des_path, 'w') as out_f:
            out_f.setparams((1, 2, Config["rate"], 0, 'NONE', 'Uncompressed'))
            out_f.writeframes(data)


def check_boundary_sample(root, moveto=None):
    def check_boundary_sample_(src):
        with wave.open(src, 'r') as f:
            cnt = f.getnframes() // CHUNK
            rms_ary = np.zeros(cnt)
            for i in range(cnt):
                f.setpos(CHUNK * i)
                chunk_data = f.readframes(CHUNK)
                rms_ary[i] = audioop.rms(chunk_data, 2)

        msk = rms_ary > 1000
        return rms_ary.max(initial=0), rms_ary[0], rms_ary[-1], rms_ary.sum(), msk.sum()

    temp_file = "../data/temp.wav"
    files = glob(root + filters[0])
    for file in files:
        with wave.open(file, 'r') as wf:
            data = wf.readframes(wf.getnframes())

        wav_data = np.frombuffer(data, dtype=np.int16)
        wav_data = np.float32(wav_data)
        mm = max(abs(wav_data)) + 1
        wav_data = wav_data * 32000 / mm
        wav_data = np.int16(wav_data)
        data = wav_data.tobytes()
        with wave.open(temp_file, 'w') as out_f:
            out_f.setparams((1, 2, Config["rate"], 0, 'NONE', 'Uncompressed'))
            out_f.writeframes(data)

        rms_max, rms_start, rms_end, rms_sum, msk_sum = check_boundary_sample_(temp_file)
        # if rms_start > 10000 or rms_end > 10000:
        # if rms_sum < 100000:
        # if msk_sum < 3:
        if rms_start < 1000 and rms_end < 1000:
            print(file, rms_max, rms_start, rms_end, rms_sum)
            if moveto:
                name = os.path.basename(file)
                des = os.path.join(moveto, name)
                shutil.move(file, des)

    os.remove(temp_file)


def move_file(file_list, des_):
    with open(file_list) as f:
        error_files = f.readlines()
    error_files = [c.strip() for c in error_files]

    for f in error_files:
        name = os.path.basename(f)
        des = os.path.join(des_, name)
        shutil.move(f, des)


def check_same(root):
    def get_file_list():
        list1 = []
        list2 = []

        list2 += glob(root + "/**/*.wav", recursive=True)

        for f in tqdm(list2):
            if (-1 != f.find('record_')) or \
                    (-1 != f.find('old_')) or \
                    (-1 != f.find('crop_')):
                list1.append(f)

        for f in tqdm(list1):
            list2.remove(f)

        return list1, list2

    def get_checksums(file_list):
        def get_checksum(file):
            with wave.open(file, 'r') as wf:
                data = wf.readframes(wf.getnframes())
            data = np.frombuffer(data, dtype=np.int16)
            ret = np.sum(abs(data))
            return ret

        checksums = []
        for f in tqdm(file_list):
            checksums.append(get_checksum(f))
        return checksums

    def get_similarity(f1, f2):
        with wave.open(f1, 'r') as wf:
            data1 = wf.readframes(wf.getnframes())
        with wave.open(f2, 'r') as wf:
            data2 = wf.readframes(wf.getnframes())
        data1 = np.frombuffer(data1, dtype=np.int16)
        data2 = np.frombuffer(data2, dtype=np.int16)
        ret = np.sum(abs(data1 - data2))
        return ret

    def move_same():
        def move_same1(file1, file2):
            file1 = file1.strip()
            file2 = file2.strip()
            file1_dir = os.path.dirname(file1)
            file1_name = os.path.basename(file1)
            file2_dir = os.path.dirname(file2)
            file2_name = os.path.basename(file2)

            des = os.path.join("../data/temp", file1_name)

            if os.path.exists(file1):
                shutil.move(file1, des)
                if file1_dir != file2_dir:
                    if os.path.exists(file2):
                        des = os.path.join(file1_dir, file2_name)
                        shutil.move(file2, des)
            else:
                if os.path.exists(file2):
                    des = os.path.join("../data/temp", file2_name)
                    shutil.move(file2, des)

        with open("../data/same.txt") as same_lines:
            for same_line in same_lines:
                same_line = same_line.strip()
                line = same_line.split(',')
                move_same1(line[0], line[1])

    files1, files2 = get_file_list()
    checksum1 = get_checksums(files1)
    checksum2 = get_checksums(files2)
    with open("../data/same.txt", 'w') as ft:
        for i in tqdm(range(len(checksum1))):
            for j in (range(len(checksum2))):
                if checksum1[i] == checksum2[j]:
                    similarity = get_similarity(files1[i], files2[j])
                    if similarity < 1000000:
                        # print(files1[i], " , ", files2[j], " , ", str(similarity))
                        ft.write(files1[i])
                        ft.write(" , ")
                        ft.write(files2[j])
                        ft.write(" , ")
                        ft.write(str(similarity))
                        ft.write("\n")

    exit(0)
    move_same()


def check_same1(root):
    def get_checksums(file_list):
        def get_checksum(file):
            with wave.open(file, 'r') as wf:
                data = wf.readframes(wf.getnframes())
            data = np.frombuffer(data, dtype=np.int16)
            ret = np.sum(abs(data))
            return ret

        checksums = []
        for f in tqdm(file_list):
            checksums.append(get_checksum(f))
        return checksums

    def get_similarity(f1, f2):
        with wave.open(f1, 'r') as wf:
            data1 = wf.readframes(wf.getnframes())
        with wave.open(f2, 'r') as wf:
            data2 = wf.readframes(wf.getnframes())
        data1 = np.frombuffer(data1, dtype=np.int16)
        data2 = np.frombuffer(data2, dtype=np.int16)
        ret = np.sum(abs(data1 - data2))
        return ret

    def move_same():
        def move_same1(file1, file2):
            file1 = file1.strip()
            file2 = file2.strip()
            file1_dir = os.path.dirname(file1)
            file1_name = os.path.basename(file1)
            file2_dir = os.path.dirname(file2)
            file2_name = os.path.basename(file2)

            des = os.path.join("../data/temp", file1_name)

            if os.path.exists(file1):
                shutil.move(file1, des)
                if file1_dir != file2_dir:
                    if os.path.exists(file2):
                        des = os.path.join(file1_dir, file2_name)
                        shutil.move(file2, des)
            else:
                if os.path.exists(file2):
                    des = os.path.join("../data/temp", file2_name)
                    shutil.move(file2, des)

        with open("../data/same.txt") as same_lines:
            for same_line in same_lines:
                same_line = same_line.strip()
                line = same_line.split(',')
                move_same1(line[0], line[1])

    files = glob(root + "/**/*.wav", recursive=True)
    checksum = get_checksums(files)
    with open("../data/same.txt", 'w') as ft:
        for i in tqdm(range(len(checksum))):
            for j in (range(i+1, len(checksum))):
                if checksum[i] == checksum[j]:
                    similarity = get_similarity(files[i], files[j])
                    if similarity < 1000000:
                        # print(files1[i], " , ", files2[j], " , ", str(similarity))
                        ft.write(files[i])
                        ft.write(" , ")
                        ft.write(files[j])
                        ft.write(" , ")
                        ft.write(str(similarity))
                        ft.write("\n")

    exit(0)
    move_same()


def annotation(root, ann_txt):
    def annotation1(path, target: str, times=1):
        files = glob(path + filters[0])
        for file in files:
            for i in range(times):
                f.write(file + ',' + target + '\n')

    with open(ann_txt, 'w') as f:
        subs = glob(root + "/*")
        for sub in subs:
            annotation1(sub + "/±", '0.5')
            annotation1(sub + "/+", '1')
            annotation1(sub + "/-", '0')
            annotation1(sub + "/±/repeat", '0.5', 4)
            annotation1(sub + "/+/repeat", '1', 4)
            annotation1(sub + "/-/repeat", '0', 4)
            annotation1(sub + "/±/err", '0.5', 8)
            annotation1(sub + "/+/err", '1', 8)
            annotation1(sub + "/-/err", '0', 8)


def check_filter(root, flt, moveto=None):
    files = glob(root + flt)
    print(root, flt, len(files))
    if moveto:
        for file in files:
            print(file)
            name = os.path.basename(file)
            des = os.path.join(moveto, name)
            shutil.move(file, des)


def check_G_sample(moveto=None):
    def check_G_text(g_text):
        def num_ok(item):
            now = item[0][0]
            for i in item[1:]:
                now += 4
                if now != i[0]:
                    return False
            return True

        def cls_ok(item):
            seq = ['', '010', '0', '1', '01', '10']
            s = ""
            now = -1
            for i in item:
                if i[1] != -1:
                    if now != i[1]:
                        s += str(i[1])
                    now = i[1]
            return s in seq, s

        def move_sub():
            dirs = ['/-', '/+', '/TBD']
            for i in items:
                basename = line[0] + "_%02d.wav" % i[0]
                dir_ = "../data/train/G" + dirs[i[1]]
                src = os.path.join(dir_, basename)
                des = os.path.join(moveto, basename)
                if os.path.exists(src):
                    shutil.move(src, des)
                # if os.path.exists(des):
                #   shutil.move(des, src)

        with open(g_text) as lines:
            for line in lines:
                line = line.strip()
                line = line.split(',')
                if len(line) > 1:
                    items = []
                    for it in line[1:]:
                        num_, cls_ = it[1:-1].split(' ')
                        items.append((int(num_), int(cls_)))
                    ok, sq = cls_ok(items)
                    if not num_ok(items) or not ok:
                        print(line, sq)
                        if moveto:
                            move_sub()

    def get_sample_info(f):
        dir_ = os.path.dirname(f)
        (shot_name, extension) = os.path.splitext(os.path.basename(f))
        name_, num_ = shot_name.split('_')[0:2]
        if dir_.find("-") != -1:
            cls_ = 0
        elif dir_.find("+") != -1:
            cls_ = 1
        else:
            cls_ = -1
        return name_.strip(), int(num_), cls_

    files = glob("../data/train/G/**" + filters[0], recursive=True)
    files = sorted(files, key=os.path.basename)

    with open("../data/G.txt", 'w') as f_out:
        name = ""
        for file in tqdm(files):
            fn, num, cls = get_sample_info(file)
            if name != fn:
                f_out.write("\n" + fn)
                name = fn
            f_out.write(",(%d %d)" % (num, cls))

    check_G_text("../data/G.txt")
    exit(0)


if __name__ == '__main__':
    # mp3_to_uniwave("../data/all")
    # wave_to_uniwave("../data/赵佳朋")
    # classify_uniwave("../data/uniwave", "../data/uniwave/-wave")
    # unniwave_crop_0sample("../data/uniwave")
    # unniwave_crop_1sample("../data/uniwave/+wave")
    # unniwave_reco_1sample("../data/uniwave/+wave")

    # check_mute_sample("../data/train/crop/TBD", flag='all', moveto="../data/temp")
    # check_loud_sample("../data/train/-", "../data/temp")

    # move_file('./1.txt', "../data/train/+/repeat")
    # move_file('./0.txt', "../data/train/+/repeat")

    # check_same1("../data/train/负样本")
    # check_by_time("../data/train/+", filters[0])
    # rename_by_time("../data/old/-", filters[0])
    # rename_by_filter("../data/train", filters[9], "crop_")

    # check_filter("../data/train/20210607/TBD", filters[0])
    # check_G_sample()
    # check_G_sample("../data/temp")

    # check_boundary_sample("../data/train/20220316/TBD/111", "../data/temp")
    # normalize_sample("../data/train/+/mute", "../data/temp")

    # annotation("../data/train")

    print('Finished!')
